# 调用预训练resnet
from keras.applications.resnet50 import ResNet50
from keras.preprocessing import image
from keras.applications.resnet50 import preprocess_input, decode_predictions
from keras.layers import Input, GlobalAveragePooling2D, Dense
from keras.models import Model
import numpy as np
import os


def train():


    # 构建不带分类器的预训练模型
    base_model = ResNet50(weights='imagenet', include_top=False, input_shape=(512,512,3))

    # 添加平均池化层
    x = base_model.output
    x = GlobalAveragePooling2D()(x)

    # 添加一个全连接层层
    x = Dense(1024, activation='relu')(x)

    # 添加一个分类器， 假设我们有5类
    predictions = Dense(5, activation='softmax', name='outclass')(x)

    # 构建我们需要训练的完整模型
    model = Model(inputs=base_model.input, outputs=predictions)

    model.summary()
# # 首先， 我们只训练顶部的几层（随机初始化的层）
# # 锁住所有Resnet50的卷积层
# for layer in base_model.layers:
#     layer.trainable = False
#
# # 编译模型（一定要在锁层以后操作）
# model.compile(optimizer='sgd', loss='categorical_crossentropy')
#
# # 在新的数据集上训练几代
# model.fit_generator(...)
#
# # 现在顶层应该训练好了，让我们开始微调Resnet50的卷积层
# # 我们会锁住底下的几层， 然后训练其余的顶层
#
# # 让我们看看每一层的名字和层号，看看我们应该锁多少层呢：
    for i, layer in enumerate(base_model.layers):
        print(i, layer.name)
#
# # 我们选择训练最上面的两个Inception block
# # 也就是说锁住前面249层，然后放开之后的层
# for layer in model.layers[:249]:
#     layer.trainable = False
# for layer in model.layers[:]:
#     layer.trainable = True
#
# # 我们需要重新编译模型， 才能使上面的修改生效
# # 让我们设置一个很低的学习率，使用SGD来微调
#
# from keras.optimizers import SGD
# model.compile(optimizer=SGD(lr=0.0001, momentum=0.9), loss='categorical_crossentropy')
#
# # 我们继续训练模型， 这次我们训练最后两个 Inception block
# # 和两个全连接层
# model.fit_generator()


if __name__ == '__main__':
    os.environ["CUDA_VISIBLE_DEVICES"] = "0"
    train()




